#!/usr/bin/env python3
"""
AIQuant系统健康检查脚本

这个脚本用于检查AIQuant系统各个组件的健康状态，包括：
- 数据库连接状态
- Redis连接状态
- API服务状态
- Web服务状态
- 系统资源使用情况
- 关键服务进程状态
"""

import os
import sys
import asyncio
import argparse
import logging
import json
import time
from pathlib import Path
from typing import Dict, Any, Optional
import psutil
import requests
import asyncpg
import redis

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from aiquant.config.settings import get_settings

logger = logging.getLogger(__name__)

class HealthChecker:
    """系统健康检查器"""
    
    def __init__(self):
        self.settings = get_settings()
        self.results = {}
        
    async def check_database(self) -> Dict[str, Any]:
        """检查数据库连接"""
        result = {
            "service": "database",
            "status": "unknown",
            "details": {},
            "timestamp": time.time()
        }
        
        try:
            # 测试连接
            conn = await asyncpg.connect(
                host=self.settings.db_host,
                port=self.settings.db_port,
                user=self.settings.db_user,
                password=self.settings.db_password,
                database=self.settings.db_name,
                timeout=5
            )
            
            # 执行简单查询
            version = await conn.fetchval("SELECT version();")
            table_count = await conn.fetchval("""
                SELECT COUNT(*) FROM information_schema.tables 
                WHERE table_schema = 'public';
            """)
            
            await conn.close()
            
            result.update({
                "status": "healthy",
                "details": {
                    "host": self.settings.db_host,
                    "port": self.settings.db_port,
                    "database": self.settings.db_name,
                    "version": version.split()[1] if version else "unknown",
                    "table_count": table_count,
                    "response_time_ms": round((time.time() - result["timestamp"]) * 1000, 2)
                }
            })
            
        except Exception as e:
            result.update({
                "status": "unhealthy",
                "error": str(e),
                "details": {
                    "host": self.settings.db_host,
                    "port": self.settings.db_port,
                    "database": self.settings.db_name
                }
            })
            
        return result
        
    def check_redis(self) -> Dict[str, Any]:
        """检查Redis连接"""
        result = {
            "service": "redis",
            "status": "unknown",
            "details": {},
            "timestamp": time.time()
        }
        
        try:
            # 创建Redis连接
            r = redis.Redis(
                host=self.settings.redis_host,
                port=self.settings.redis_port,
                password=self.settings.redis_password or None,
                db=self.settings.redis_db,
                socket_timeout=5,
                socket_connect_timeout=5
            )
            
            # 测试连接
            r.ping()
            
            # 获取Redis信息
            info = r.info()
            
            result.update({
                "status": "healthy",
                "details": {
                    "host": self.settings.redis_host,
                    "port": self.settings.redis_port,
                    "version": info.get("redis_version", "unknown"),
                    "used_memory": info.get("used_memory_human", "unknown"),
                    "connected_clients": info.get("connected_clients", 0),
                    "response_time_ms": round((time.time() - result["timestamp"]) * 1000, 2)
                }
            })
            
        except Exception as e:
            result.update({
                "status": "unhealthy",
                "error": str(e),
                "details": {
                    "host": self.settings.redis_host,
                    "port": self.settings.redis_port
                }
            })
            
        return result
        
    def check_api_service(self) -> Dict[str, Any]:
        """检查API服务"""
        result = {
            "service": "api",
            "status": "unknown",
            "details": {},
            "timestamp": time.time()
        }
        
        try:
            # 检查健康端点
            url = f"http://{self.settings.api_host}:{self.settings.api_port}/health"
            response = requests.get(url, timeout=10)
            
            if response.status_code == 200:
                health_data = response.json()
                result.update({
                    "status": "healthy",
                    "details": {
                        "url": url,
                        "status_code": response.status_code,
                        "response_time_ms": round((time.time() - result["timestamp"]) * 1000, 2),
                        "health_data": health_data
                    }
                })
            else:
                result.update({
                    "status": "unhealthy",
                    "error": f"HTTP {response.status_code}",
                    "details": {
                        "url": url,
                        "status_code": response.status_code
                    }
                })
                
        except Exception as e:
            result.update({
                "status": "unhealthy",
                "error": str(e),
                "details": {
                    "url": f"http://{self.settings.api_host}:{self.settings.api_port}/health"
                }
            })
            
        return result
        
    def check_web_service(self) -> Dict[str, Any]:
        """检查Web服务"""
        result = {
            "service": "web",
            "status": "unknown",
            "details": {},
            "timestamp": time.time()
        }
        
        try:
            # 检查Web服务端口
            web_host = getattr(self.settings, 'web_host', '0.0.0.0')
            web_port = getattr(self.settings, 'web_port', 8501)
            
            url = f"http://{web_host}:{web_port}"
            response = requests.get(url, timeout=10)
            
            result.update({
                "status": "healthy" if response.status_code == 200 else "degraded",
                "details": {
                    "url": url,
                    "status_code": response.status_code,
                    "response_time_ms": round((time.time() - result["timestamp"]) * 1000, 2)
                }
            })
            
        except Exception as e:
            result.update({
                "status": "unhealthy",
                "error": str(e),
                "details": {
                    "url": f"http://{getattr(self.settings, 'web_host', '0.0.0.0')}:{getattr(self.settings, 'web_port', 8501)}"
                }
            })
            
        return result
        
    def check_system_resources(self) -> Dict[str, Any]:
        """检查系统资源"""
        result = {
            "service": "system",
            "status": "unknown",
            "details": {},
            "timestamp": time.time()
        }
        
        try:
            # CPU使用率
            cpu_percent = psutil.cpu_percent(interval=1)
            
            # 内存使用情况
            memory = psutil.virtual_memory()
            
            # 磁盘使用情况
            disk = psutil.disk_usage('/')
            
            # 网络连接数
            connections = len(psutil.net_connections())
            
            # 进程数
            process_count = len(psutil.pids())
            
            # 负载平均值（仅Unix系统）
            load_avg = None
            try:
                load_avg = os.getloadavg()
            except (OSError, AttributeError):
                pass
                
            # 判断系统状态
            status = "healthy"
            warnings = []
            
            if cpu_percent > 80:
                status = "degraded"
                warnings.append(f"CPU使用率过高: {cpu_percent}%")
                
            if memory.percent > 85:
                status = "degraded"
                warnings.append(f"内存使用率过高: {memory.percent}%")
                
            if disk.percent > 90:
                status = "degraded"
                warnings.append(f"磁盘使用率过高: {disk.percent}%")
                
            result.update({
                "status": status,
                "details": {
                    "cpu": {
                        "usage_percent": cpu_percent,
                        "count": psutil.cpu_count()
                    },
                    "memory": {
                        "total_gb": round(memory.total / (1024**3), 2),
                        "used_gb": round(memory.used / (1024**3), 2),
                        "available_gb": round(memory.available / (1024**3), 2),
                        "usage_percent": memory.percent
                    },
                    "disk": {
                        "total_gb": round(disk.total / (1024**3), 2),
                        "used_gb": round(disk.used / (1024**3), 2),
                        "free_gb": round(disk.free / (1024**3), 2),
                        "usage_percent": disk.percent
                    },
                    "network": {
                        "connections": connections
                    },
                    "processes": {
                        "count": process_count
                    },
                    "load_average": load_avg,
                    "warnings": warnings
                }
            })
            
        except Exception as e:
            result.update({
                "status": "unhealthy",
                "error": str(e)
            })
            
        return result
        
    def check_processes(self) -> Dict[str, Any]:
        """检查关键进程"""
        result = {
            "service": "processes",
            "status": "unknown",
            "details": {},
            "timestamp": time.time()
        }
        
        try:
            # 要检查的进程关键词
            target_processes = [
                "uvicorn",
                "streamlit",
                "celery",
                "python"
            ]
            
            running_processes = []
            all_processes = []
            
            for proc in psutil.process_iter(['pid', 'name', 'cmdline', 'cpu_percent', 'memory_percent']):
                try:
                    proc_info = proc.info
                    cmdline = ' '.join(proc_info['cmdline']) if proc_info['cmdline'] else ''
                    
                    # 检查是否是目标进程
                    for target in target_processes:
                        if target in proc_info['name'].lower() or target in cmdline.lower():
                            if 'aiquant' in cmdline.lower():
                                running_processes.append({
                                    "pid": proc_info['pid'],
                                    "name": proc_info['name'],
                                    "cmdline": cmdline[:100] + "..." if len(cmdline) > 100 else cmdline,
                                    "cpu_percent": proc_info['cpu_percent'],
                                    "memory_percent": round(proc_info['memory_percent'], 2)
                                })
                                
                    all_processes.append(proc_info['name'])
                    
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    continue
                    
            # 判断状态
            status = "healthy" if running_processes else "degraded"
            
            result.update({
                "status": status,
                "details": {
                    "aiquant_processes": running_processes,
                    "total_processes": len(all_processes),
                    "process_count_by_name": {name: all_processes.count(name) 
                                           for name in set(all_processes) 
                                           if all_processes.count(name) > 1}
                }
            })
            
        except Exception as e:
            result.update({
                "status": "unhealthy",
                "error": str(e)
            })
            
        return result
        
    async def run_all_checks(self) -> Dict[str, Any]:
        """运行所有健康检查"""
        print("🔍 开始系统健康检查...")
        
        checks = [
            ("数据库", self.check_database()),
            ("Redis", self.check_redis()),
            ("API服务", self.check_api_service()),
            ("Web服务", self.check_web_service()),
            ("系统资源", self.check_system_resources()),
            ("进程状态", self.check_processes())
        ]
        
        results = {}
        overall_status = "healthy"
        
        for name, check_coro in checks:
            print(f"  检查 {name}...")
            
            if asyncio.iscoroutine(check_coro):
                result = await check_coro
            else:
                result = check_coro
                
            results[result["service"]] = result
            
            # 更新整体状态
            if result["status"] == "unhealthy":
                overall_status = "unhealthy"
            elif result["status"] == "degraded" and overall_status == "healthy":
                overall_status = "degraded"
                
        # 汇总结果
        summary = {
            "overall_status": overall_status,
            "timestamp": time.time(),
            "checks": results,
            "summary": {
                "total_checks": len(results),
                "healthy": sum(1 for r in results.values() if r["status"] == "healthy"),
                "degraded": sum(1 for r in results.values() if r["status"] == "degraded"),
                "unhealthy": sum(1 for r in results.values() if r["status"] == "unhealthy")
            }
        }
        
        return summary
        
    def print_results(self, results: Dict[str, Any]):
        """打印检查结果"""
        status_icons = {
            "healthy": "✅",
            "degraded": "⚠️",
            "unhealthy": "❌",
            "unknown": "❓"
        }
        
        print(f"\n{'='*60}")
        print(f"AIQuant 系统健康检查报告")
        print(f"{'='*60}")
        print(f"整体状态: {status_icons.get(results['overall_status'], '❓')} {results['overall_status'].upper()}")
        print(f"检查时间: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(results['timestamp']))}")
        print(f"检查项目: {results['summary']['total_checks']} 项")
        print(f"  - 健康: {results['summary']['healthy']} 项")
        print(f"  - 降级: {results['summary']['degraded']} 项")
        print(f"  - 异常: {results['summary']['unhealthy']} 项")
        
        print(f"\n{'='*60}")
        print("详细检查结果:")
        print(f"{'='*60}")
        
        for service, result in results["checks"].items():
            icon = status_icons.get(result["status"], "❓")
            print(f"\n{icon} {service.upper()}: {result['status']}")
            
            if "error" in result:
                print(f"   错误: {result['error']}")
                
            if "details" in result and result["details"]:
                for key, value in result["details"].items():
                    if isinstance(value, dict):
                        print(f"   {key}:")
                        for k, v in value.items():
                            print(f"     {k}: {v}")
                    elif isinstance(value, list) and value:
                        print(f"   {key}: {len(value)} 项")
                        for item in value[:3]:  # 只显示前3项
                            if isinstance(item, dict):
                                print(f"     - {item}")
                            else:
                                print(f"     - {item}")
                        if len(value) > 3:
                            print(f"     ... 还有 {len(value) - 3} 项")
                    else:
                        print(f"   {key}: {value}")


async def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="AIQuant系统健康检查脚本")
    parser.add_argument(
        "--output", "-o",
        choices=["console", "json", "file"],
        default="console",
        help="输出格式"
    )
    parser.add_argument(
        "--file", "-f",
        help="输出文件路径（当output=file时使用）"
    )
    parser.add_argument(
        "--continuous", "-c",
        type=int,
        help="持续监控模式，指定检查间隔（秒）"
    )
    parser.add_argument(
        "--alert-threshold",
        choices=["healthy", "degraded", "unhealthy"],
        default="unhealthy",
        help="告警阈值"
    )
    
    args = parser.parse_args()
    
    # 设置日志
    logging.basicConfig(
        level=logging.WARNING,  # 减少日志输出
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
    )
    
    checker = HealthChecker()
    
    try:
        if args.continuous:
            print(f"🔄 启动持续监控模式，检查间隔: {args.continuous} 秒")
            print("按 Ctrl+C 停止监控")
            
            while True:
                results = await checker.run_all_checks()
                
                if args.output == "console":
                    checker.print_results(results)
                elif args.output == "json":
                    print(json.dumps(results, indent=2, ensure_ascii=False))
                elif args.output == "file":
                    output_file = args.file or f"health_check_{int(time.time())}.json"
                    with open(output_file, 'w', encoding='utf-8') as f:
                        json.dump(results, f, indent=2, ensure_ascii=False)
                    print(f"结果已保存到: {output_file}")
                    
                # 检查是否需要告警
                alert_levels = {"healthy": 0, "degraded": 1, "unhealthy": 2}
                if alert_levels.get(results["overall_status"], 2) >= alert_levels.get(args.alert_threshold, 2):
                    print(f"\n🚨 系统状态达到告警阈值: {results['overall_status']}")
                    
                print(f"\n等待 {args.continuous} 秒后进行下次检查...")
                await asyncio.sleep(args.continuous)
                
        else:
            # 单次检查
            results = await checker.run_all_checks()
            
            if args.output == "console":
                checker.print_results(results)
            elif args.output == "json":
                print(json.dumps(results, indent=2, ensure_ascii=False))
            elif args.output == "file":
                output_file = args.file or f"health_check_{int(time.time())}.json"
                with open(output_file, 'w', encoding='utf-8') as f:
                    json.dump(results, f, indent=2, ensure_ascii=False)
                print(f"结果已保存到: {output_file}")
                
            # 根据结果设置退出码
            exit_codes = {"healthy": 0, "degraded": 1, "unhealthy": 2, "unknown": 3}
            sys.exit(exit_codes.get(results["overall_status"], 3))
            
    except KeyboardInterrupt:
        print("\n👋 健康检查已停止")
    except Exception as e:
        print(f"\n❌ 健康检查失败: {e}")
        sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())